# coding:utf-8
# -*- coding: utf-8 -*-
# @Time : 2020/11/6 15:51
# @Author: 王鹏-wangpeng99
# @File : MyTree.py

import networkx as nx
from treelib import Node, Tree, plugins
import matplotlib.pyplot as plt
import random, time


class MyTree(Tree):
    def __init__(self, tree=None, deep=False, node_class=None, identifier=None):
        super(MyTree, self).__init__(tree, deep, node_class, identifier)
        # 记录子树信息
        self._subTrees = []
        # 画图板
        self._g = nx.Graph()
        # matplotlib支持中文
        plt.rcParams['font.sans-serif'] = ['KaiTi']  # 指定默认字体
        plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题
        plt.rcParams['savefig.dpi'] = 300  # 图片像素
        plt.rcParams['figure.dpi'] = 300  # 分辨率

    def paste(self, nid, new_tree, deep=False):
        #先存储这个子树
        self._subTrees.append(new_tree)
        # 调用Tree的方法
        super(MyTree, self).paste(nid, new_tree, deep)

    # 构造子树
    def buildSubTree(self, basenode, treeNameBase=None, nodeNumber=10):
        # assert isinstance(basenode, (Tree,MyTree))
        if basenode is None:
            raise OSError("basenode can't be None")
        if not self.contains(basenode):
            raise plugins.NodeIDAbsentError("Node '%s' is not in the tree" % basenode)
        if not treeNameBase:
            raise ValueError('没有给定treeBase')
        subtree = Tree()
        subtree.create_node(treeNameBase, treeNameBase)

        random.seed(int(round(time.time() * 1000000)))
        tmplst = [treeNameBase]
        for i in range(nodeNumber):
            node = '{}{}'.format(treeNameBase, i)
            subtree.create_node(node, node, parent=random.choice(tmplst))
            tmplst.append(node)
        self.paste(basenode, subtree)

    # 递归获取所有的节点关联关系
    def __parseNodeLink(self, node=None, nodelinklst=None):
        children = self.children(node)
        if children and len(children) > 0:
            for child in children:
                nodelinklst.append((node, child.identifier))
                self.__parseNodeLink(child.identifier, nodelinklst)
        else:
            return

    def drawTree(self):
        # 获取节点信息
        allNodes = self.all_nodes()
        nodeInfolst = [(item.identifier, {'tag': item.tag, 'identifier': item.identifier, 'data': item.data}) for item
                       in allNodes]
        nodeIdentifiers = [item[0] for item in nodeInfolst]
        # 获取边信息
        nodelinklst = []
        self.__parseNodeLink(self.root, nodelinklst)
        self._g = nx.Graph()
        # 增加节点信息
        self._g.add_nodes_from(nodeInfolst)
        # 增加边信息
        self._g.add_edges_from(nodelinklst)
        # 构造pos信息
        # pos = nx.fruchterman_reingold_layout(self._g, iterations=100)
        pos = nx.spring_layout(self._g)
        pos['root'] = (0.0, 0.0)
        ##############################
        ############ 画图开始##########
        ##############################
        # 先画根节点
        nx.draw_networkx_nodes(self._g, pos=pos,nodelist=['root'], node_shape='*', node_color='k')
        # 画节点标签
        nx.draw_networkx_labels(self._g, pos=pos, labels={'root': 'root'}, font_color='r', font_size='10')
        #画root到第一个子节点的边
        # rootChildren=self.children(self.root)[0]
        # print(rootChildren)
        # 画边信息
        nx.draw_networkx_edges(self._g, pos=pos, edgelist=[(self.root,self.children(self.root)[0].identifier)], edge_color='k')
        # 再画各个子树：
        for i in range(len(self._subTrees)):
            # 开始获取子树基本信息,浅拷贝子树
            subtree =Tree(self._subTrees[i])
            subtree.show()
            subtreeAllNodes = subtree.all_nodes()
            subtreeNodeIdentifiers = [item.identifier for item in subtreeAllNodes]
            # 子树的边信息
            subtreeNodeLinklst = []
            self.__parseNodeLink(subtree.root, subtreeNodeLinklst)
            print(subtreeNodeLinklst)
            # 先画子树的根节点
            nx.draw_networkx_nodes(self._g, pos=pos, nodelist=[subtreeNodeIdentifiers[0]], node_shape='*', node_color=[plt.cm.tab10(i)])
            # 画节点标签
            nx.draw_networkx_labels(self._g, pos=pos, labels={subtreeNodeIdentifiers[0]: subtreeNodeIdentifiers[0]}, font_color='k', font_size='8')
            # 再画子树的子节点
            # 确定各个子树的边颜色
            nodecolor = [plt.cm.tab10(i) for j in range(len(subtreeNodeIdentifiers[1:]))]
            edgecolor = [plt.cm.tab10(i+5) for j in range(len(subtreeNodeIdentifiers[1:]))]
            nx.draw(self._g, pos=pos,
                    # 子树节点信息
                    nodelist=subtreeNodeIdentifiers[1:], node_shape='.', node_color=nodecolor,  # node_size=800,
                    # 子树边信息
                    edgelist=subtreeNodeLinklst, edge_color=edgecolor,
                    # 子树标签
                    labels={item: item for item in subtreeNodeIdentifiers[1:]}, font_color='k', font_size='8')
        plt.show()


if __name__ == '__main__':
    tree = MyTree()
    tree.create_node("root", "root")  # root node
    centerlst = ['综拓中心', '主数据中心', '智能中心', '再保中心', '用户中心', '销售中心', '理赔中心', '客户中心', '健康管理中心', '工具平台', '费用中心', '单证中心',
                 '承保中心', '产品中心', '财务中心']
    centerchoice = random.sample(centerlst, k=5)
    for center in centerchoice:
        if len(tree) == 1:
            tree.buildSubTree('root', center)
        else:
            basenode = random.choice(tree.leaves())
            tree.buildSubTree(basenode.identifier, center)

    tree.show()
    tree.drawTree()
    # for item in tree._subTrees:
    #     item.show()
    # tree.drawTree()
